#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import argparse
import pandas as pd


def merge_files(input_path, output_file, file_type='csv', blacklist=[]):
    # Validate if the input path exists
    if not os.path.exists(input_path):
        return "Input path does not exist"

    # Obtain absolute path
    input_path = os.path.abspath(input_path)

    # If output file is not provided, create a default one based on input path
    if not output_file:
        output_dir = os.path.basename(input_path)
        output_file = os.path.join(input_path, f"{output_dir}.{file_type}")
    else:
        # Make sure the provided output file ends with the specified file type
        if not output_file.endswith(f".{file_type}"):
            return f"Output file must have a '{file_type}' extension"
        output_file = os.path.abspath(output_file)

    # Append input_path to each file in the blacklist
    blacklist_with_path = [os.path.join(input_path, file) for file in blacklist]

    # Get all files of the specified type in the input path, excluding those in the blacklist
    files_to_merge = []
    for root, dirs, files in os.walk(input_path):
        for file in files:
            file_path = os.path.join(root, file)
            # Exclude output files and files in the blacklist
            if file_path != output_file and file_path.endswith(
                    f".{file_type}") and file_path not in blacklist_with_path:
                files_to_merge.append(file_path)

    # Check if there are any files to merge
    if not files_to_merge:
        return f"No {file_type} files found in the path {input_path}"

    # Create an empty DataFrame to store the merged data
    merged_data = None
    first_file = True

    # Merge files one by one
    for file_to_merge in files_to_merge:
        try:
            if file_type == 'xlsx':
                data = pd.read_excel(file_to_merge, engine='openpyxl')
            elif file_type == 'csv':
                # Read CSV with auto-detected separator (, or \t)
                data = pd.read_csv(file_to_merge, sep=None, engine='python')
            elif file_type == 'txt':
                with open(file_to_merge, 'r', encoding='utf-8') as txt_file:
                    first_line = txt_file.readline()
                    if '\t' in first_line:
                        sep = '\t'
                    else:
                        sep = ','
                    # Return to the beginning of the file
                    txt_file.seek(0)
                    data = pd.read_csv(txt_file, sep=sep)
            else:
                return f"Unsupported file type: {file_type}"

            # If it's the first file, assign it to merged_data directly
            if first_file:
                merged_data = data
                first_file = False
            else:
                # Otherwise, attempt to merge the data with merged_data
                merged_data = pd.concat([merged_data, data], ignore_index=True, sort=False)

        except Exception as e:
            return f"Error merging file {file_to_merge}: {str(e)}"

        # If no data was successfully merged, report an error
        if merged_data is None:
            return f"No files with matching columns found in the path {input_path}"

    # Save the merged data to the output file
    try:
        if file_type == 'xlsx':
            merged_data.to_excel(output_file, index=False)
        else:
            # Both CSV and TXT file types are allowed, so we use 'to_csv' for both
            merged_data.to_csv(output_file, index=False)
        return f"Merge successful, output file name: {output_file}"
    except Exception as e:
        return f"Error saving merged data: {str(e)}"


def main():
    parser = argparse.ArgumentParser(description="Merge files of a specified type")
    parser.add_argument("input_path", help="Path to the files to merge")
    parser.add_argument("-o", "--output_file",
                        help="Output file name, default is the last directory name of the input path plus the file type")
    parser.add_argument("-f", "--file_type", default="csv", choices=["csv", "xlsx", "txt"],
                        help="File type to merge, default is csv")
    parser.add_argument("-b", "--blacklist", nargs="*", default=[], help="Files to exclude from merging (must be within input_path)")

    args = parser.parse_args()

    result = merge_files(args.input_path, args.output_file, args.file_type, args.blacklist)
    print(result)


if __name__ == '__main__':
    main()
